package com.example.sharding.jdbc.demo.util;

import java.time.Duration;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.*;

/**
 * 按月分片工具类
 *
 * @author airtrioa
 * @date 2021/4/13
 */
public class MonthPreciseShardingUtil {
  protected static final List<String> TABLES = new ArrayList<>();
  private static final String TEST_LOG = "test_log";
  private static final String TEST_RECORD = "test_record";

  static {
    TABLES.add(TEST_LOG);
    TABLES.add(TEST_RECORD);
  }

  private static final Integer ONE = 1;
  private static final Integer TWELVE = 12;
  private static final Integer THREE_HUNDRED_SIXTY_FIVE = 365;

  private MonthPreciseShardingUtil() {
  }

  /**
   * 处理表名。仅获取表名中的序列号
   *
   * @param tableName 表名
   * @return 序列号Str
   */
  public static String obtainTableIndex(String tableName) {
    for (String table : TABLES) {
      if (tableName.startsWith(table)) {
        return tableName.replace(table, "");
      }
    }
    throw new IllegalArgumentException(String.format("暂不支持此表[%s]的分表查询", tableName));
  }

  /**
   * 根据时间上下界，过滤相应的表名
   *
   * @param originalTables 表名列表
   * @param lowerEndpoint  下界
   * @param upperEndpoint  上界
   * @return 表名列表
   */
  public static Set<String> rangeShardingTables(Collection<String> originalTables, LocalDateTime lowerEndpoint, LocalDateTime upperEndpoint) {
    Set<String> tables = new HashSet<>();
    // 获取时间范围内的 月份集合
    Set<Integer> mouthSet = MonthPreciseShardingUtil.getSuffixListForRange(lowerEndpoint, upperEndpoint);
    for (String tableName : originalTables) {
      // 获取表的序列号
      String tableIndex = MonthPreciseShardingUtil.obtainTableIndex(tableName);
      if (mouthSet.contains(Integer.parseInt(tableIndex))) {
        // 若月份中包含表的序列，则把表名放进集合
        tables.add(tableName);
      }
    }
    return tables;
  }

  /**
   * 根据返回，获取到范围涉及的月份
   *
   * @param lowerEndpoint 下限
   * @param upperEndpoint 上限
   * @return 月份集合
   */
  public static Set<Integer> getSuffixListForRange(LocalDateTime lowerEndpoint, LocalDateTime upperEndpoint) {
    Set<Integer> sets = new TreeSet<>();
    if (Duration.between(lowerEndpoint, upperEndpoint).toDays() > THREE_HUNDRED_SIXTY_FIVE) {
      // 时间跨度大于一年，直接返回12个月
      for (int i = ONE; i <= TWELVE; i++) {
        sets.add(i);
      }
      return sets;
    }
    int lowerEndpointMonth = lowerEndpoint.getMonth().getValue();
    int upperEndpointMonth = upperEndpoint.getMonth().getValue();

    if (lowerEndpointMonth <= upperEndpointMonth) {
      for (int i = lowerEndpointMonth; i <= upperEndpointMonth; i++) {
        sets.add(i);
      }
    } else {
      for (int i = ONE; i <= upperEndpointMonth; i++) {
        sets.add(i);
      }
      for (int i = lowerEndpointMonth; i <= TWELVE; i++) {
        sets.add(i);
      }
    }
    return sets;
  }

  /**
   * 根据返回，获取到范围涉及的月份
   *
   * @param lowerEndpoint 下限
   * @param upperEndpoint 上限
   * @return 月份集合
   */
  public static Set<Integer> getSuffixListForRange(LocalDate lowerEndpoint, LocalDate upperEndpoint) {
    return getSuffixListForRange(LocalDateTime.of(lowerEndpoint, LocalTime.MIN), LocalDateTime.of(upperEndpoint, LocalTime.MAX));
  }
}
